package org.wenshu.ai;

import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.jina.JinaEmbeddingModel;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.ChatMessage;

import com.fasterxml.jackson.databind.ObjectMapper;

import org.junit.jupiter.api.Test;

import java.util.*;
import java.util.function.Function;

import static java.time.Duration.ofSeconds;

/**
 * Java实现的多智能体聊天系统
 * @author AI Generated
 */
class JavaMultiAgentIT {

    /**
     * DeepSeek智能体类
     */
    static class DeepSeekAgent {
        private final String name;
        private final String systemPrompt;
        private final int shortMemSize;
        private final List<Map.Entry<String, String>> shortMem = new ArrayList<>(); // 保存最近的对话
        private final List<Map.Entry<String, Embedding>> longMem = new ArrayList<>(); // 存储记忆向量
        private final OpenAiChatModel model;
        private final EmbeddingModel embeddingModel;

        /**
         * 构造函数
         * @param name 智能体名称
         * @param systemPrompt 系统提示
         * @param shortMemSize 短期记忆大小
         */
        public DeepSeekAgent(String name, String systemPrompt, int shortMemSize) {
            this.name = name;
            this.systemPrompt = systemPrompt != null ? systemPrompt : "You are a helpful assistant";
            this.shortMemSize = shortMemSize;
            
            // 配置OpenAI客户端
            this.model = OpenAiChatModel.builder()
                    .baseUrl("https://api.deepseek.com")
                    .apiKey("your-key")
                    .modelName("deepseek-chat")
                    .logRequests(true)
                    .logResponses(true)
                    .build();
            
            // 配置嵌入模型
            this.embeddingModel = JinaEmbeddingModel.builder()
                    .apiKey("your-key")
                    .modelName("jina-embeddings-v3")
                    .timeout(ofSeconds(10))
                    .logRequests(true)
                    .logResponses(true)
                    .build();
        }

        /**
         * 构建提示信息
         * @param conversationHistory 对话历史
         * @return 构建的提示信息
         */
        private String buildPrompt(List<Map<String, String>> conversationHistory) {
            StringBuilder prompt = new StringBuilder();
            
            // 添加系统提示
            prompt.append(this.systemPrompt).append("\n");
            
            // 构建最近的对话历史
            prompt.append("[对话历史]\n");
            int startIdx = Math.max(0, conversationHistory.size() - this.shortMemSize);
            for (int i = startIdx; i < conversationHistory.size(); i++) {
                Map<String, String> entry = conversationHistory.get(i);
                prompt.append(entry.get("speaker")).append(": ").append(entry.get("content")).append("\n");
            }
            
            // 添加相关记忆（如果有的话）
            if (!longMem.isEmpty()) {
                try {
                    // 获取当前话题
                    String currentTopic = conversationHistory.get(conversationHistory.size() - 1).get("content");
                    
                    // 将当前话题转换为嵌入向量
                    Embedding currentVec = embeddingModel.embed(currentTopic).content();
                    
                    // 计算相似度并找到最匹配的记忆
                    double maxScore = -1;
                    int bestMatchIndex = -1;
                    
                    for (int i = 0; i < longMem.size(); i++) {
                        Embedding memVec = longMem.get(i).getValue();
                        double score = cosineSimilarity(currentVec.vector(), memVec.vector());
                        
                        if (score > maxScore) {
                            maxScore = score;
                            bestMatchIndex = i;
                        }
                    }
                    
                    if (bestMatchIndex != -1) {
                        String bestMatch = longMem.get(bestMatchIndex).getKey();
                        prompt.append("\n[相关记忆]\n").append(bestMatch).append("\n");
                    }
                } catch (Exception e) {
                    // 如果嵌入过程出错，忽略相关记忆部分
                    System.err.println("查找相关记忆时出错: " + e.getMessage());
                }
            }
            
            // 添加回复指示
            prompt.append("\n请回应 ")
                  .append(conversationHistory.get(conversationHistory.size() - 1).get("speaker"))
                  .append(" 并分享您的观点...");
            
            return prompt.toString();
        }

        /**
         * 计算余弦相似度
         * @param vec1 向量1
         * @param vec2 向量2
         * @return 余弦相似度
         */
        private double cosineSimilarity(float[] vec1, float[] vec2) {
            double dotProduct = 0.0;
            double norm1 = 0.0;
            double norm2 = 0.0;
            
            for (int i = 0; i < vec1.length; i++) {
                dotProduct += vec1[i] * vec2[i];
                norm1 += Math.pow(vec1[i], 2);
                norm2 += Math.pow(vec2[i], 2);
            }
            
            return dotProduct / (Math.sqrt(norm1) * Math.sqrt(norm2));
        }

        /**
         * 响应对话历史
         * @param conversationHistory 对话历史
         * @return 智能体的响应
         */
        public String respond(List<Map<String, String>> conversationHistory) {
            // 构建提示
            String prompt = buildPrompt(conversationHistory);
            
            // 获取LLM响应
            List<ChatMessage> messages = Arrays.asList(
                new SystemMessage(this.systemPrompt),
                new UserMessage(prompt)
            );
            
            String response = model.generate(messages).content().text();
            
            // 更新记忆
            updateMemory(conversationHistory.get(conversationHistory.size() - 1).get("content"), response);
            
            return response;
        }

        /**
         * 更新记忆
         * @param lastQuestion 最后的问题
         * @param response 响应
         */
        private void updateMemory(String lastQuestion, String response) {
            // 更新短期记忆
            shortMem.add(Map.entry(lastQuestion, response));
            if (shortMem.size() > shortMemSize) {
                shortMem.remove(0);
            }
            
            // 更新长期记忆
            String memoryText = "Q: " + lastQuestion + "\nA: " + response;
            try {
                Embedding memoryVector = embeddingModel.embed(memoryText).content();
                longMem.add(Map.entry(memoryText, memoryVector));
            } catch (Exception e) {
                System.err.println("更新长期记忆时出错: " + e.getMessage());
            }
        }

        /**
         * 获取智能体名称
         * @return 智能体名称
         */
        public String getName() {
            return name;
        }
    }

    /**
     * 多智能体聊天
     * @param agents 智能体列表
     * @param rounds 对话轮数
     * @param initialPrompt 初始提示
     */
    private static void multiAgentChat(List<DeepSeekAgent> agents, int rounds, String initialPrompt) {
        List<Map<String, String>> history = new ArrayList<>();
        
        // 添加初始提示
        Map<String, String> initialMessage = new HashMap<>();
        initialMessage.put("speaker", "Moderator");
        initialMessage.put("content", initialPrompt);
        history.add(initialMessage);
        
        System.out.println("Moderator: " + initialPrompt);
        System.out.println("-".repeat(50));
        
        // 进行多轮对话
        for (int i = 0; i < rounds; i++) {
            DeepSeekAgent currentAgent = agents.get(i % agents.size());
            System.out.println("\nCurrent Speaker: " + currentAgent.getName());
            
            String response = currentAgent.respond(history);
            
            Map<String, String> message = new HashMap<>();
            message.put("speaker", currentAgent.getName());
            message.put("content", response);
            history.add(message);
            
            System.out.println(currentAgent.getName() + ": " + response);
            System.out.println("-".repeat(50));
        }
    }

    @Test
    void testMultiAgentChat() {
        // 初始化三个州的智能体，并设置短期记忆大小
        DeepSeekAgent florida = new DeepSeekAgent("佛罗里达",
                "你代表佛罗里达州，这是一个关键的摇摆州，拥有大量老年人和拉丁裔人口。你需要回应其他州的观点，并从佛罗里达州的角度分析美国大选。",
                2);  // Florida记住最近2轮对话
        
        DeepSeekAgent pennsylvania = new DeepSeekAgent("宾夕法尼亚",
                "你代表宾夕法尼亚州，这是一个工业州，蓝领工人和郊区选民对选举结果至关重要。你需要回应其他州的观点，并从宾夕法尼亚州的角度分析美国大选。",
                3);  // Pennsylvania记住最近3轮对话
        
        DeepSeekAgent michigan = new DeepSeekAgent("密歇根",
                "你代表密歇根州，这是一个以汽车工业闻名的州。工会和少数族裔选民对选举有重要影响。你需要回应其他州的观点，并从密歇根州的角度分析美国大选。",
                4);  // Michigan记住最近4轮对话
        
        // 进行5轮对话
        multiAgentChat(
                Arrays.asList(florida, pennsylvania, michigan),
                5,
                "2024年美国大选即将来临。你认为哪个党派将在摇摆州占据优势？请讨论特朗普和民主党。模拟这次选举。"
        );
    }

    /**
     * 交互式测试（可选）
     */
    public static void main(String[] args) {
        // 初始化智能体
        DeepSeekAgent florida = new DeepSeekAgent("Florida",
                "You represent Florida, a key swing state with a large elderly and Latino population. You need to respond to the views of other states and analyze the U.S. election from Florida's perspective.",
                2);
        
        DeepSeekAgent pennsylvania = new DeepSeekAgent("Pennsylvania",
                "You represent Pennsylvania, an industrial state where blue-collar workers and suburban voters are crucial to the election outcome. You need to respond to the views of other states and analyze the U.S. election from Pennsylvania's perspective.",
                3);
        
        DeepSeekAgent michigan = new DeepSeekAgent("Michigan",
                "You represent Michigan, a state known for its automotive industry. Union and minority voters have a significant impact on the election. You need to respond to the views of other states and analyze the U.S. election from Michigan's perspective.",
                4);
        
        List<DeepSeekAgent> agents = Arrays.asList(florida, pennsylvania, michigan);
        
        // 创建对话历史
        List<Map<String, String>> history = new ArrayList<>();
        
        // 添加初始提示
        Scanner scanner = new Scanner(System.in);
        System.out.print("请输入初始提示: ");
        String initialPrompt = scanner.nextLine();
        
        Map<String, String> initialMessage = new HashMap<>();
        initialMessage.put("speaker", "主持人");
        initialMessage.put("content", initialPrompt);
        history.add(initialMessage);
        
        System.out.println("主持人: " + initialPrompt);
        System.out.println("-".repeat(50));
        
        // 交互式对话
        int round = 0;
        while (true) {
            DeepSeekAgent currentAgent = agents.get(round % agents.size());
            System.out.println("\n当前发言者: " + currentAgent.getName());
            
            String response = currentAgent.respond(history);
            
            Map<String, String> message = new HashMap<>();
            message.put("speaker", currentAgent.getName());
            message.put("content", response);
            history.add(message);
            
            System.out.println(currentAgent.getName() + ": " + response);
            System.out.println("-".repeat(50));
            
            System.out.print("继续对话? (y/n): ");
            String answer = scanner.nextLine();
            if ("n".equalsIgnoreCase(answer)) {
                break;
            }
            round++;
        }
        
        scanner.close();
    }}